from typing import TypedDict, Optional, List
from langgraph.graph import StateGraph, END
from IPython.display import display


# TypedDict‌：TypedDict通过显式声明键的类型和可选性，提供了严格的字典结构。
# 它允许开发者定义严格的字典结构，并在编码阶段通过工具（如 mypy ）进行静态类型检查，
# 从而在编码阶段捕获类型错误‌


class AgentState(TypedDict):
    user_input: str              # 用户输入
    user_type: Optional[str]     # "newbie" or "experienced"
    conversation_history: List[str]  # 对话历史
    current_step: str            # 当前处于哪个对话阶段

## 每经过一个节点（Node），都会对 state 进行更新，比如判断用户是不是新手、是否注册成功等。

## **state 是什么意思？
## ** 这是 Python 的 字典解包（dictionary unpacking） 语法
## 把原来的 state 字典里的所有键值对复制过来，然后再添加或覆盖一个新的键 "current_step"，值为 "greeting"。

## 为什么这么做？
# 在 LangGraph / LangChain 构建的流程中，每个节点接收一个 state 对象，并返回一个新的 state 对象。你不应该直接修改原始的 state（即不可变原则），而是每次返回一个更新后的副本。
##
def greeting(state: AgentState) -> AgentState:
    print("机器人：您好，我是OKX官方助手Qwen，很高兴为您服务。您是第一次接触加密货币交易吗？")
    return {**state, "current_step": "greeting"}


def determine_user_type(state: AgentState) -> AgentState:
    user_input = state["user_input"].lower()

    if any(word in user_input for word in ["是", "第一次", "没经验"]):
        print("机器人：没关系，我们可以从基础开始。首先，我来简单介绍一下我们的平台...")
        return {**state, "user_type": "newbie", "current_step": "introduce_newbie"}
    else:
        print("机器人：明白了，那我可以先了解一下您平时主要关注哪些方面？比如手续费政策、高级交易工具等。")
        return {**state, "user_type": "experienced", "current_step": "introduce_experienced"}

def introduce_newbie(state: AgentState) -> AgentState:
        print("机器人：在OKX上交易有三大优势：安全可靠、操作简单、币种丰富。现在注册还可以领取价值$10的比特币奖励哦～")
        print("👉 点击下方链接立即注册：[注册链接]")
        return {**state, "current_step": "finish"}

def introduce_experienced(state: AgentState) -> AgentState:
    print("机器人：作为老手，您可能会对以下功能感兴趣：流动性强、手续费折扣、专业交易工具等。")
    print("需要我帮您对比一下OKX与其他平台的区别吗？")
    return {**state, "current_step": "finish"}

def finish(state: AgentState) -> AgentState:
    print("机器人：如果您有任何疑问，欢迎随时回来找我。祝您交易顺利！")
    return {**state}


# 创建图
# Our graph can now handle two key tasks:
#
# Each node can receive the current State as input and output an update to the state.
# Updates to messages will be appended to the existing list rather than overwriting it, thanks to the prebuilt add_messages function used with the Annotated syntax.
#
workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("greeting", greeting)
workflow.add_node("determine_user_type", determine_user_type)
workflow.add_node("introduce_newbie", introduce_newbie)
workflow.add_node("introduce_experienced", introduce_experienced)
workflow.add_node("finish", finish)

# 设置入口点
workflow.set_entry_point("greeting")

# 添加边
workflow.add_edge("greeting", "determine_user_type")
workflow.add_conditional_edges(
    "determine_user_type",
    ## 如果 state["user_type"] == "newbie"，则跳转到 "introduce_newbie" 节点；
    ## 如果 state["user_type"] == "experienced"，则跳转到 "introduce_experienced" 节点。
    lambda state: state["user_type"],
    {
        "newbie": "introduce_newbie",
        "experienced": "introduce_experienced"
    }
)
workflow.add_edge("introduce_newbie", "finish")
workflow.add_edge("introduce_experienced", "finish")
workflow.add_edge("finish", END)

## 运行
app = workflow.compile()
# display(app.get_graph().draw_mermaid_png())
print(app.get_graph().draw_mermaid())
initial_state = {
    "user_input": "是的，我第一次接触。",
    "conversation_history": [],
    "current_step": ""
}

result = app.invoke(initial_state)